#include "yuv2rgb.h"
#include <stdio.h>

#define BLOCK_W 16
#define BLOCK_H 16

// 支持 yuvj420p 格式的核函数
__global__ void kernel_yuv2rgb(unsigned char *buf_y, unsigned char *buf_uv, int stride_y, int stride_uv, float *rgb, int width, int height, bool is_jpeg_range) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x < width && y < height) {
        // 根据格式选择不同的处理方式
        float fy, fu, fv;
        
        // 读取 Y 分量
        fy = (float)buf_y[y * stride_y + x];
        
        // 读取 UV 分量 (NV12 格式)
        int uv_x = x / 2;
        int uv_y = y / 2;
        int uv_index = uv_y * stride_uv + uv_x * 2;
        fu = (float)buf_uv[uv_index] - 128.0f;
        fv = (float)buf_uv[uv_index + 1] - 128.0f;

        // 计算 RGB
        float r, g, b;
        
        if (is_jpeg_range) {
            // JPEG 范围的转换公式 (全范围 0-255)
            r = fy + 1.402f * fv;
            g = fy - 0.344136f * fu - 0.714136f * fv;
            b = fy + 1.772f * fu;
        } else {
            // TV 范围的转换公式 (有限范围 16-235)
            fy = (fy - 16.0f) * 255.0f / 219.0f;  // 扩展 Y 到全范围
            r = fy + 1.402f * fv;
            g = fy - 0.344136f * fu - 0.714136f * fv;
            b = fy + 1.772f * fu;
        }
        
        // 限制范围在 [0, 255] 然后归一化到 [0, 1]
        r = fminf(fmaxf(r, 0.0f), 255.0f) / 255.0f;
        g = fminf(fmaxf(g, 0.0f), 255.0f) / 255.0f;
        b = fminf(fmaxf(b, 0.0f), 255.0f) / 255.0f;
        
        // 修改数据布局：使用交错格式而不是平面格式
        int rgb_index = (y * width + x) * 3;
        rgb[rgb_index] = r;     // R
        rgb[rgb_index + 1] = g; // G  
        rgb[rgb_index + 2] = b; // B
    }
}

// 现有的 copy_rgb_rect 函数保持不变
__global__ void kernel_copy_rgb_rect(float *src, float *dst, rect_t rect, int width, int height, float scale) {
    int dst_width  = rect.w;
    int dst_height = rect.h;
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x < dst_width && y < dst_height) {
        int i_src_y = (int)(y * scale + rect.y + 0.5f);
        int i_src_x = (int)(x * scale + rect.x + 0.5f);
        if (i_src_x >= width)
            i_src_x = width - 1;
        if (i_src_y >= height)
            i_src_y = height - 1;
        int i_src = (int)(i_src_y * width + i_src_x);
        int i_dst = y * dst_width + x;
        float *src_r = src;
        float *src_g = src + width * height;
        float *src_b = src + width * height * 2;
        float *dst_r = dst;
        float *dst_g = dst + dst_width * dst_height;
        float *dst_b = dst + dst_width * dst_height * 2;
        dst_r[i_dst] = src_r[i_src];
        dst_g[i_dst] = src_g[i_src];
        dst_b[i_dst] = src_b[i_src];
    }
}

// 修改 yuv2rgb 函数签名
extern "C" void yuv2rgb(unsigned char *y, unsigned char *uv, int stride_y, int stride_uv, float *rgb, int width, int height, bool is_jpeg_range) {
    dim3 block_size      = { BLOCK_W, BLOCK_H, 1 };
    dim3 grid_size       = { ((unsigned int)width + BLOCK_W - 1) / BLOCK_W, ((unsigned int)height + BLOCK_H - 1) / BLOCK_H, 1 };

    kernel_yuv2rgb<<<grid_size, block_size>>>(y, uv, stride_y, stride_uv, rgb, width, height, is_jpeg_range);
    cudaDeviceSynchronize();
}

extern "C" void copy_rgb_rect(float *src, float *dst, rect_t *rect, int width, int height, float scale) {
    dim3 block_size      = { BLOCK_W, BLOCK_H, 1 };
    dim3 grid_size       = { ((unsigned int)rect->w + BLOCK_W - 1) / BLOCK_W, ((unsigned int)rect->h + BLOCK_H - 1) / BLOCK_H, 1 };

    kernel_copy_rgb_rect<<<grid_size, block_size>>>(src, dst, *rect, width, height, scale);
    cudaDeviceSynchronize();
}
